#include "afpb_service_generator.h"

#include <google/protobuf/descriptor.h>
#include <google/protobuf/io/printer.h>
#include <google/protobuf/stubs/strutil.h>

#include "google/protobuf/compiler/cpp/cpp_helpers.h"
#include "google/protobuf/compiler/cpp/cpp_options.h"
#include "internal_options.pb.h"

AfpbServiceGenerator::AfpbServiceGenerator(
    const pb::ServiceDescriptor *descriptor,
    const pbcpp::Options &options)
  : descriptor_(descriptor) {
  vars_["classname"] = descriptor_->name();
  vars_["full_name"] = descriptor_->full_name();
  if ( options.dllexport_decl.empty() ) {
    vars_["dllexport"] = "";
  } else {
    vars_["dllexport"] = options.dllexport_decl + " ";
  }

  if ( descriptor_->options().HasExtension(
           asfinger::protocol::service_options) ) {
    vars_["original_hash"] = "  typedef std::integral_constant<uint32_t, 0x" +
        pb::ToUpper(pb::ToHex(HashServiceName(
            descriptor_->options()
                .GetExtension(asfinger::protocol::service_options)
                .descriptor_name()))) +
        "u> OriginalHash;\n";
  } else {
    vars_["original_hash"] = "";
  }
  vars_["name_hash"] = "  typedef std::integral_constant<uint32_t, 0x" +
      pb::ToUpper(pb::ToHex(HashServiceName(descriptor_->full_name()))) +
      "u> NameHash;\n";
}

AfpbServiceGenerator::~AfpbServiceGenerator() = default;

void AfpbServiceGenerator::GenerateDeclarations(pb::io::Printer *printer) {
  GenerateInterface(printer);
}

void AfpbServiceGenerator::GenerateInterface(pb::io::Printer *printer) {
  printer->Print(vars_,
                 "class $dllexport$$classname$\n"
                 "    : public ServiceBase {\n"
                 " public:\n"
                 "\n"
                 "  explict $classname$(bool use_original_hash);\n"
                 "  virtual ~$classname$;\n"
                 "$original_hash$"
                 "$name_hash$");

  printer->Indent();

  printer->Print(
      vars_,
      "\n"
      "static ::google::protobuf::ServiceDescriptor const* descriptor();\n"
      "\n");

  if ( !descriptor_->options().HasExtension(
           asfinger::protocol::sdk_service_options) ||
       descriptor_->options()
           .GetExtension(asfinger::protocol::sdk_service_options)
           .inbound() ) {
    printer->Print(
        vars_,
        "// client methods --------------------------------------------------\n");

    GenerateClientMethodSignatures(printer);

    printer->Print(vars_, "\n");
  }

  printer->Print(
      "void CallServerMethod(uint32_t token, uint32_t method_id, MessageBuffer buffer) final;\n");

  if ( !descriptor_->options().HasExtension(
           asfinger::protocol::sdk_service_options) ||
       descriptor_->options()
           .GetExtension(asfinger::protocol::sdk_service_options)
           .outbound() ) {
    printer->Outdent();

    printer->Print(
        "\n"
        " protected:\n");

    printer->Indent();

    printer->Print(
        vars_,
        "// server methods --------------------------------------------------\n");

    GenerateServerMethodSignatures(printer);
  }

  printer->Outdent();

  printer->Print(vars_,
                 "\n"
                 " private:\n"
                 "  uint32_t service_hash_{0};\n"
                 "\n"
                 "  GOOGLE_DISALLOW_EVIL_CONSTRUCTORS($classname$);\n"
                 "};\n");
}

void AfpbServiceGenerator::GenerateImplementation(pb::io::Printer *printer) {
  printer->Print(
      vars_,
      "$classname$::$classname$(bool use_original_hash)\n"
      "    : service_hash_(use_original_hash ? OriginalHash::value : NameHash::value) {\n"
      "}\n"
      "\n"
      "$classname$::~$classname$() {\n"
      "}\n"
      "\n"
      "::google::protobuf::ServiceDescriptor const* $classname$::descriptor() {\n"
      "  protobuf_AssignDescriptorsOnce();\n"
      "  return $classname$_descriptor_;\n"
      "}\n"
      "\n");

  if ( !descriptor_->options().HasExtension(
           asfinger::protocol::sdk_service_options) ||
       descriptor_->options()
           .GetExtension(asfinger::protocol::sdk_service_options)
           .inbound() ) {
    GenerateClientMethodImplementations(printer);
  }

  if ( !descriptor_->options().HasExtension(
           asfinger::protocol::sdk_service_options) ||
       descriptor_->options()
           .GetExtension(asfinger::protocol::sdk_service_options)
           .outbound() ) {
    GenerateServerCallMethod(printer);
    GenerateServerImplementations(printer);
  } else {
    printer->Print(
        vars_,
        "void $classname$::CallServerMethod(uint32_t token, uint32_t method_id, MessageBuffer /*buffer*/) {\n"
        "  LOG_ERROR(\"{} Server tried to call server method {}\",\n"
        "      GetCallerInfo().c_str(), methodId);\n"
        "  LOG_ERROR(\"service.protobuf\", \"%s Server tried to call server method %u\",\n"
        "    GetCallerInfo().c_str(), method_id);\n"
        "}\n"
        "\n");
  }
}

void AfpbServiceGenerator::GenerateDescriptorInitializer(
    pb::io::Printer *printer,
    int index) {
  std::map<std::string, std::string> vars;
  vars["classname"] = descriptor_->name();
  vars["index"]     = pb::SimpleItoa(index);

  printer->Print(vars, "$classname$_descriptor_ = file->service($index$);\n");
}

void AfpbServiceGenerator::GenerateClientMethodSignatures(
    pb::io::Printer *printer) {
  for ( int i = 0; i < descriptor_->method_count(); ++i ) {
    const auto *method = descriptor_->method(i);
    if ( !method->options().HasExtension(asfinger::protocol::method_options) ) {
      continue;
    }
    std::map<std::string, std::string> sub_vars;
    sub_vars["name"]      = method->name();
    sub_vars["full_name"] = descriptor_->name() + "." + method->name();
    sub_vars["method_id"] =
        pb::SimpleItoa(method->options()
                           .GetExtension(asfinger::protocol::method_options)
                           .id());
    sub_vars["input_type"]      = pbcpp::ClassName(method->input_type(), true);
    sub_vars["output_type"]     = pbcpp::ClassName(method->output_type(), true);
    sub_vars["input_type_name"] = method->input_type()->full_name();

    if ( "NoResponse" != method->output_type()->name() ) {
      printer->Print(
          sub_vars,
          "void $name$(const $input_type$ *request, std::function<void(const $output_type$ *)> response_callback, bool client = false, bool server = false);\n");
    } else {
      printer->Print(
          sub_vars,
          "void $name$(const $input_type$ *request, bool client = false, bool server = false);\n");
    }
  }
}

void AfpbServiceGenerator::GenerateServerMethodSignatures(
    pb::io::Printer *printer) {
  for ( int i = 0; i < descriptor_->method_count(); ++i ) {
    const auto *method = descriptor_->method(i);
    if ( !method->options().HasExtension(asfinger::protocol::method_options) ) {
      continue;
    }
    std::map<std::string, std::string> sub_vars;
    sub_vars["name"]        = method->name();
    sub_vars["input_type"]  = pbcpp::ClassName(method->input_type(), true);
    sub_vars["output_type"] = pbcpp::ClassName(method->output_type(), true);

    if ( "NoResponse" != method->output_type()->name() ) {
      printer->Print(
          sub_vars,
          "virtual uint32_t Handle$name$(const $input_type$ *request, $output_type$ *response, std::function<void(ServiceBase *, uint32_t, const ::google::protobuf::Message *)> &continuation);\n");
    } else {
      printer->Print(
          sub_vars,
          "virtual uint32_t Handle$name$(const $input_type$ *request);\n");
    }
  }
}

void AfpbServiceGenerator::GenerateClientMethodImplementations(
    pb::io::Printer *printer) {
  for ( int i = 0; i < descriptor_->method_count(); ++i ) {
    const auto *method = descriptor_->method(i);
    if ( !method->options().HasExtension(asfinger::protocol::method_options) ) {
      continue;
    }

    std::map<std::string, std::string> sub_vars;
    sub_vars["classname"] = vars_["classname"];
    sub_vars["name"]      = method->name();
    sub_vars["full_name"] = descriptor_->name() + "." + method->name();
    sub_vars["method_id"] =
        pb::SimpleItoa(method->options()
                           .GetExtension(asfinger::protocol::method_options)
                           .id());
    sub_vars["input_type"]      = pbcpp::ClassName(method->input_type(), true);
    sub_vars["output_type"]     = pbcpp::ClassName(method->output_type(), true);
    sub_vars["input_type_name"] = method->input_type()->full_name();

    if ( "NoResponse" != method->output_type()->name() ) {
      printer->Print(
          sub_vars,
          "void $classname$::$name$(const $input_type$ *request, std::function<void(const $output_type$ *)> response_callback, bool client /*= false*/, bool server /*= false*/) {\n"
          "  LOG_DEBUG(\"{} Server called client method $full_name$($input_type_name$( {} ))\",\n"
          "      GetCallerInfo().c_str(), request->ShortDebugString().c_str());\n"
          "  std::function<void(MessageBuffer> callback = [response_callback](MessageBuffer buffer) -> void {\n"
          "    $output_type$ response;\n"
          "    if (response.ParseFromArray(buffer.GetReadPointer(), buffer.GetActiveSize())) {\n"
          "      response_callback(&response);\n"
          "    }\n"
          "  };\n"
          "  SendRequest(service_hash_, $method_id$ | (client ? 0x40000000 : 0) | (server ? 0x80000000 : 0), request, std::move(callback));\n"
          "}\n"
          "\n");

    } else {
      printer->Print(
          sub_vars,
          "void $classname$::$name$(const $input_type$ *request, bool client = false, bool server = false) {\n"
          "  LOG_DEBUG(\"%s Server called client method $full_name$($input_type_name$( {} ))\",\n"
          "      GetCallerInfo().c_str(), request->ShortDebugString().c_str());\n"
          "  SendRequest(service_hash_, $method_id$ | (client ? 0x40000000 : 0) | (server ? 0x80000000 : 0), request);\n"
          "}\n"
          "\n");
    }
  }
}

void AfpbServiceGenerator::GenerateServerCallMethod(pb::io::Printer *printer) {
  printer->Print(
      vars_,
      "void $classname$::CallServerMethod(uint32_t token, uint32_t method_id, MessageBuffer buffer) {\n"
      "  switch(methodId & 0x3FFFFFFF) {\n");

  for ( int i = 0; i < descriptor_->method_count(); ++i ) {
    const auto *method = descriptor_->method(i);
    if ( !method->options().HasExtension(asfinger::protocol::method_options) ) {
      continue;
    }

    std::map<std::string, std::string> sub_vars;
    sub_vars["classname"] = vars_["classname"];
    sub_vars["name"]      = method->name();
    sub_vars["full_name"] = descriptor_->name() + "." + method->name();
    sub_vars["method_id"] =
        pb::SimpleItoa(method->options()
                           .GetExtension(asfinger::protocol::method_options)
                           .id());
    sub_vars["input_type"]      = pbcpp::ClassName(method->input_type(), true);
    sub_vars["output_type"]     = pbcpp::ClassName(method->output_type(), true);
    sub_vars["input_type_name"] = method->input_type()->full_name();
    sub_vars["output_type_name"] = method->output_type()->full_name();

    printer->Print(
        sub_vars,
        "  case $method_id$:\n"
        "  {\n"
        "    $input_type$ request;\n"
        "    if (!request.ParseFromArray(buffer.GetReadPointer(), buffer.GetActiveSize())) {\n"
        "      LOG_DEBUG(\"{} Failed to parse request for $full_name$ server method call.\", GetCallerInfo().c_str());\n"
        "      SendResponse(service_hash_, method_id, token, kRpcErrorCode_Malformed_Request);\n"
        "      return;\n"
        "    }\n");

    if ( "NoResponse" != method->output_type()->name() ) {
      printer->Print(
          sub_vars,
          "    LOG_DEBUG(\"{} Client called server method $full_name$($input_type_name$( {} )).\",\n"
          "      GetCallerInfo().c_str(), request->ShortDebugString().c_str());\n"
          "    std::function<void(ServiceBase *, uint32_t, const ::google::protobuf::Message *)> continuation = [token, method_id](ServiceBase* service, uint32_t status, const ::google::protobuf::Message *response) {\n"
          "      DBG_ASSERT(response->GetDescriptor() == $output_type$::descriptor());\n"
          "      $classname$* self = static_cast<$classname$*>(service);\n"
          "      LOG_DEBUG(\"{} Client called server method $full_name$() returned $output_type_name$( {} ) status {}.\",\n"
          "        self->GetCallerInfo().c_str(), response->ShortDebugString().c_str(), status);\n"
          "      if (!status) {\n"
          "        self->SendResponse(self->service_hash_, method_id, token, response);\n"
          "      } else {\n"
          "        self->SendResponse(self->service_hash_, method_id, token, status);\n"
          "      }\n"
          "    }\n"
          "    $output_type$ response;\n"
          "    uint32_t status = Handle$name$(&request, &response, continuation);\n"
          "    if (continuation) {\n"
          "      continuation(this, status, &response);\n"
          "    }\n");
    } else {
      printer->Print(
          sub_vars,
          "    uint32_t status = Handle$name$(&request);\n"
          "    LOG_DEBUG(\"{} Client called server method $full_name$() returned $output_type_name$( {} ) status {}.\",\n"
          "      self->GetCallerInfo().c_str(), response->ShortDebugString().c_str(), status);\n"
          "    if (status) {\n"
          "      SendResponse(service_hash_, methodId, token, status);\n"
          "    }\n");
    };

    printer->Print(sub_vars,
                   "    break;\n"
                   "  }\n");
  }

  printer->Print(
      vars_,
      "  default:\n"
      "    TC_LOG_ERROR(\"Bad method id {}.\", method_id);\n"
      "    SendResponse(service_hash_, method_id, token, kRpcErrorCode_Invalid_Method);\n"
      "    break;\n"
      "  }\n"
      "}\n"
      "\n");
}

void AfpbServiceGenerator::GenerateServerImplementations(
    pb::io::Printer *printer) {
  for ( int i = 0; i < descriptor_->method_count(); ++i ) {
    const auto *method = descriptor_->method(i);
    if ( !method->options().HasExtension(asfinger::protocol::method_options) ) {
      continue;
    }

    std::map<std::string, std::string> sub_vars;
    sub_vars["classname"]   = vars_["classname"];
    sub_vars["name"]        = method->name();
    sub_vars["full_name"]   = descriptor_->name() + "." + method->name();
    sub_vars["input_type"]  = pbcpp::ClassName(method->input_type(), true);
    sub_vars["output_type"] = pbcpp::ClassName(method->output_type(), true);

    if ( "NoResponse" != method->output_type()->name() ) {
      printer->Print(
          sub_vars,
          "uint32_t $classname$::Handle$name$(const $input_type$ *request, $output_type$ *response, std::function<void(ServiceBase *, uint32_t, const ::google::protobuf::Message *)> &continuation) {\n"
          "  LOG_ERROR(\"{} Client tried to call not implemented method $full_name$( {} )\",\n"
          "    GetCallerInfo().c_str(), request->ShortDebugString().c_str());\n"
          "  return kRpcErrorCode_Not_Implemented;\n"
          "}\n"
          "\n");
    } else {
      printer->Print(
          sub_vars,
          "uint32_t $classname$::Handle$name$(const $input_type$ *request) {\n"
          "  LOG_ERROR(\"{} Client tried to call not implemented method $full_name$( {} )\",\n"
          "    GetCallerInfo().c_str(), request->ShortDebugString().c_str());\n"
          "  return kRpcErrorCode_Not_Implemented;\n"
          "}\n"
          "\n");
    };
  }
}

std::uint32_t AfpbServiceGenerator::HashServiceName(const std::string &name) {
  std::uint32_t hash = 1315423911;

  for ( auto &&c : name ) {
    hash ^= ((hash << 5) + c + (hash >> 2));
  }

  return hash & 0x7FFFFFFF;
}
